package main.java.lpr;

import java.awt.image.BufferedImage;
import java.util.ArrayList;
import java.util.List;

import org.bytedeco.javacv.Frame;
import org.bytedeco.javacv.Java2DFrameConverter;
import org.bytedeco.javacv.OpenCVFrameConverter;
import org.bytedeco.opencv.opencv_core.Mat;
import org.bytedeco.opencv.opencv_core.Point;
import org.bytedeco.opencv.opencv_core.Point2f;
import org.bytedeco.opencv.opencv_core.Rect;
import org.bytedeco.opencv.opencv_core.Scalar;
import org.bytedeco.opencv.opencv_core.Size;


import ai.djl.inference.Predictor;
import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.ImageFactory;
import ai.djl.modality.cv.output.BoundingBox;
import ai.djl.modality.cv.output.DetectedObjects;
import ai.djl.modality.cv.output.Rectangle;
import ai.djl.modality.cv.output.DetectedObjects.DetectedObject;
import ai.djl.modality.cv.translator.YoloV5Translator;
import ai.djl.repository.zoo.Criteria;
import ai.djl.repository.zoo.ModelZoo;
import ai.djl.repository.zoo.ZooModel;
import ai.djl.translate.Translator;

import static org.bytedeco.opencv.global.opencv_imgproc.FONT_HERSHEY_PLAIN;
import static org.bytedeco.opencv.global.opencv_imgproc.LINE_8;
import static org.bytedeco.opencv.global.opencv_imgproc.getRectSubPix;
import static org.bytedeco.opencv.global.opencv_imgproc.putText;
import static org.bytedeco.opencv.global.opencv_imgproc.rectangle;
import static org.bytedeco.opencv.global.opencv_imgproc.resize;

/**
 * PlateAreaDetection
 * 
 * @date: 2022-02-21 16:25
 * @author: alvin
 * @copyright: Copyright © 2022 YZZH. All rights reserved.
 * @description: 基于yolov5的车牌区域识别
 */
public class PlateAreaDetection {

    private static String model_path = "models/lpi.torchscript.pt";

    private static ZooModel<Image, DetectedObjects> model = null;

    static {
        Translator<Image, DetectedObjects> translator = YoloV5Translator.builder()
                .optSynsetArtifactName("lpc.names")
                .build();
        Criteria<Image, DetectedObjects> criteria = Criteria.builder()
                .setTypes(Image.class, DetectedObjects.class)
                .optModelUrls(model_path)
                .optModelName("lpi.torchscript.pt")
                .optTranslator(translator)
                .optEngine("PyTorch")
                .build();

        try {
            model = ModelZoo.loadModel(criteria);
            System.out.println(model);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    /**
     * 数据集对象识别 接收480*480的图片
     * 
     * @param img
     * @param isDraw
     * @return
     */
    public static List<Plate> detection(Mat img, boolean isDraw, double thresh) {
        List<Plate> plates = new ArrayList<Plate>();
        if (img == null || img.empty()) {
            return plates;
        }

        Mat src = new Mat();
        Size size = new Size(480, 480);
        resize(img, src, size);
        size.close();

        // 转换成BufferedImage
        OpenCVFrameConverter.ToMat matConv = new OpenCVFrameConverter.ToMat();
        Java2DFrameConverter biConv = new Java2DFrameConverter();
        Frame frame = matConv.convert(src);
        BufferedImage bimg = biConv.getBufferedImage(frame);

        // 释放
        frame.close();
        biConv.close();
        matConv.close();

        // 转换Image
        Image djl_img = ImageFactory.getInstance().fromImage(bimg);

        // 构建预测者
        Predictor<Image, DetectedObjects> predictor = model.newPredictor();

        try {
            // 检测出全部对象
            DetectedObjects objects = predictor.predict(djl_img);

            for (DetectedObject obj : objects.<DetectedObject>items()) {
                if (thresh < obj.getProbability()) {
                    BoundingBox bbox = obj.getBoundingBox();
                    Rectangle rectangle = bbox.getBounds();

                    int h = (int) (rectangle.getHeight() * img.arrayHeight() / 480);
                    int w = (int) (rectangle.getWidth() * img.arrayWidth() / 480);
                    int x = (int) (rectangle.getX() * img.arrayWidth() / 480);
                    int y = (int) (rectangle.getY() * img.arrayHeight() / 480);

                   //coco对象
					String type = "";
					try {
						type = obj.getClassName().trim();
					} catch (Exception e) {

					}
					Plate o = new Plate();
                    o.setType(type);
                    o.setConfidence((float)obj.getProbability());
                    o.setH(h+10);
                    o.setW(w+10);
                    o.setX(x-5);//左上角的点
                    o.setY(y-5);//左上角的点
                    
                    //截取对象
                    Mat mat = new Mat();
                    Size  si = new Size(o.getW(), o.getH());
                    Point2f p = new Point2f(o.getX() + o.getW()/2, o.getY() + o.getH()/2);
                    getRectSubPix(img, si, p, mat);
                    o.setMat(mat);
                    
                    p.close();
                    si.close();
                    
                    List<PlateChar> plateChars = PlateCharDetection.detection(mat, 0.7);
                    if (!plateChars.isEmpty()) {
                        o.setPlateResult(plateChars.get(0));
                    } 

                    if(isDraw) {
                        //在原图上画box
                        String text = String.format("%s %s: %.2f", o.getPlateResult().getPlateNumber(), o.getType(), o.getConfidence());
                        
                        Rect rect = new Rect();
                        rect.x(o.getX());
                        rect.y(o.getY());
                        rect.width(o.getW());
                        rect.height(o.getH());
                        // 画框
                        rectangle(img, rect, new Scalar(0, 255, 0, 0), 2, LINE_8, 0);// print blue rectangle
                        // 画名字
                        Point pt = new Point(rect.x(), rect.y());
                        putText(img, text, pt, FONT_HERSHEY_PLAIN, 2, new Scalar(0, 0, 255, 0));
                        
                        pt.close();
                        rect.close();
                    }
                    
                    //保存对象
                    plates.add(o);
                }

            }

        } catch (Exception e) {
            e.printStackTrace();
        }

        src.close();
        predictor.close();
        return plates;
    }

}
